# 导入必要的PyTorch库
import math

import torch
import torch.nn as nn


class MLP(nn.Module):
    """多层感知机模块"""

    def __init__(self, input_dim, hidden_dim, output_dim, dropout_rate=0.0):
        """
        初始化MLP模块
        Args:
            input_dim: 输入维度
            hidden_dim: 隐藏层维度
            output_dim: 输出维度
            dropout_rate: dropout比率
        """
        super(MLP, self).__init__()
        self.activation = nn.GELU()  # 使用GELU激活函数
        self.fc1 = nn.Linear(input_dim, hidden_dim)  # 第一个全连接层
        self.fc2 = nn.Linear(hidden_dim, output_dim)  # 第二个全连接层
        self.dropout = nn.Dropout(dropout_rate)  # Dropout层用于防止过拟合

    def forward(self, x):
        """前向传播函数"""
        x = self.fc1(x)  # 通过第一个全连接层
        x = self.activation(x)  # 应用激活函数
        x = self.dropout(x)  # 应用dropout
        x = self.fc2(x)  # 通过第二个全连接层
        x = self.dropout(x)  # 再次应用dropout
        return x


class Attention(nn.Module):
    """多头自注意力机制模块"""

    def __init__(self, hidden_size, dropout_rate):
        """
        初始化注意力模块
        Args:
            hidden_size: 隐藏层维度
            dropout_rate: dropout比率
        """
        super(Attention, self).__init__()
        self.num_attention_heads = 12  # 注意力头数量
        self.attention_head_size = int(hidden_size / self.num_attention_heads)  # 每个注意力头的维度
        self.all_head_size = self.num_attention_heads * self.attention_head_size  # 所有注意力头的总维度

        # 定义Query、Key、Value的线性变换层
        self.query = nn.Linear(hidden_size, self.all_head_size)
        self.key = nn.Linear(hidden_size, self.all_head_size)
        self.value = nn.Linear(hidden_size, self.all_head_size)

        self.out = nn.Linear(hidden_size, hidden_size)  # 输出投影层
        self.attn_dropout = nn.Dropout(dropout_rate)  # 注意力dropout
        self.proj_dropout = nn.Dropout(dropout_rate)  # 投影dropout
        self.softmax = nn.Softmax(dim=-1)  # softmax层

    def transpose_for_scores(self, x):
        """重塑输入张量用于多头注意力计算"""
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)  # 重塑张量维度
        return x.permute(0, 2, 1, 3)  # 调整维度顺序

    def forward(self, hidden_states):
        """前向传播函数"""
        # 生成query、key、value向量
        mixed_query_layer = self.query(hidden_states)  # 生成query
        mixed_key_layer = self.key(hidden_states)  # 生成key
        mixed_value_layer = self.value(hidden_states)  # 生成value

        # 重塑query、key、value的维度
        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # 计算注意力分数
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))  # 计算query和key的点积
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)  # 缩放注意力分数
        attention_probs = self.softmax(attention_scores)  # 应用softmax
        attention_probs = self.attn_dropout(attention_probs)  # 应用dropout

        # 计算注意力输出
        context_layer = torch.matmul(attention_probs, value_layer)  # 与value相乘
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()  # 调整维度顺序
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)  # 重塑维度
        attention_output = self.out(context_layer)  # 通过输出投影层
        attention_output = self.proj_dropout(attention_output)  # 应用dropout
        return attention_output


class PatchEmbed(nn.Module):
    """图像patch嵌入模块"""

    def __init__(self, input_dim, img_size, patch_size, emb_dim):
        """
        初始化patch嵌入模块
        Args:
            input_dim: 输入通道数
            img_size: 输入图像大小
            patch_size: patch大小
            emb_dim: 嵌入维度
        """
        super(PatchEmbed, self).__init__()

        self.proj_layer = nn.Conv2d(input_dim, emb_dim, patch_size, patch_size)  # 使用卷积实现patch划分
        self.num_patches = (img_size // patch_size) ** 2  # 计算patch数量
        self.pos_embs = nn.Parameter(torch.zeros(1, self.num_patches, emb_dim))  # 位置编码
        self.dropout = nn.Dropout(0.1)  # dropout层

    def forward(self, x):
        """前向传播函数"""
        x = self.proj_layer(x)  # 通过投影层
        x = x.flatten(2)  # 展平空间维度
        x = x.transpose(1, 2)  # 调整维度顺序
        x = x + self.pos_embs  # 添加位置编码
        x = self.dropout(x)  # 应用dropout
        return x


class TransformerLayer(nn.Module):
    """Transformer编码器层"""

    def __init__(self, emb_dim):
        """
        初始化Transformer层
        Args:
            emb_dim: 嵌入维度
        """
        super(TransformerLayer, self).__init__()

        mlp_size = 3072  # 论文中4.4节规定的MLP大小

        # 注意力相关层
        self.attention_norm = nn.LayerNorm(emb_dim, eps=1e-6)  # 注意力层前的归一化
        self.attention_layer = Attention(hidden_size=emb_dim, dropout_rate=0.1)  # 注意力层

        # 前馈网络相关层
        self.feedforward_norm = nn.LayerNorm(emb_dim, eps=1e-6)  # 前馈网络前的归一化
        self.mlp = MLP(emb_dim, mlp_size, emb_dim, 0.1)  # MLP层

    def forward(self, x):
        """前向传播函数"""
        mhsa_norm_out = self.attention_layer(self.attention_norm(x))  # 多头自注意力处理
        x = x + mhsa_norm_out  # 残差连接
        x_ffn = self.feedforward_norm(x)  # 前馈网络归一化
        x_mlp = self.mlp(x_ffn)  # MLP处理
        x = x + x_mlp  # 残差连接
        return x
